from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.tools import tool
from utils import websearch
from typing import Annotated
from langgraph.prebuilt import create_react_agent


@tool
def web_search(query: Annotated[str, "互联网查询内容"]):
    """通过web_search工具查询互联网上的信息，可以引用相近业务代码作为参考"""
    _rt = websearch(query)
    return _rt


_researcher_system_template = """
你是一位精通代码生成的专家，专注于根据用户的具体需求，
运用Python语法为用户量身定制相应的业务代码。用户只需提出需求,
你将构思出精准而高效的代码蓝图，无需实际执行，除了注释和代码以外不要添加其他的内容。
你可以使用以下的工具协助你完成任务:
{tools_name}

"""


class Generator:

    def __init__(self, llm):

        _tools = [
            web_search,
        ]

        _prompt = ChatPromptTemplate.from_messages(
            [
                ("system", _researcher_system_template),
                MessagesPlaceholder(variable_name="messages"),
            ]
        )

        _prompt = _prompt.partial(tools_name=",".join([_tool.name for _tool in _tools]))
        _llm_with_tools_agent = create_react_agent(llm, tools=_tools)
        self._chain = _prompt | _llm_with_tools_agent

    def __call__(self, state):
        _rt = self._chain.invoke(state)
        _messages = _rt["messages"]
        return _messages[-1]


if __name__ == "__main__":
    from langchain_openai import ChatOpenAI

    _llm = ChatOpenAI(
        api_key="ollama",
        model="qwen2.5:7b",
        base_url="http://192.168.10.11:60026/v1",
        temperature=0.7,
    )

    _generator = Generator(_llm)
    _rt = _generator(
        {
            "messages": [
                (
                    "human",
                    "创建两个对象，一个叫Animal一个叫Cat，Cat继承Animal",
                )
            ]
        }
    )

    print(_rt)
